import {useCallback, useEffect, useRef, useState} from 'react'

const baseUrl = import.meta.env.VITE_WS_BASE_URL

export const useWebsocket = (token: string | null) => {
  const ws = useRef<WebSocket | null>(null)
  const [readyState, setReadyState] = useState<number>(WebSocket.CONNECTING)

  useEffect(() => {
    if (!token) return

    const websocket = new WebSocket(baseUrl + `/infra/ws?token=${token}&source=desktop`)
    websocket.onopen = () => {
      console.log('WebSocket Connected')
      setReadyState(websocket.readyState)
    }

    websocket.onclose = () => {
      console.log('WebSocket Disconnected')
      setReadyState(websocket.readyState)
    }

    websocket.onerror = (error) => {
      console.error('WebSocket Error:', error)
    }

    websocket.onmessage = (event) => {
      const message = JSON.parse(event.data)
      console.log('Received message:', message)
    }

    ws.current = websocket

    return () => {
      websocket.close()
    }
  }, [token])

  const sendMessage = useCallback((message: string) => {
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      ws.current.send(JSON.stringify(message))
    }
  }, [])

  return {sendMessage, readyState}
}
